import cv2
import numpy as np
import onnxruntime as ort
import time
import serial

SERIAL_PORT = '/dev/serial0'
BAUD_RATE = 115200
SEND_INTERVAL = 1.0  
RESOLUTION = (320, 320)  
FRAMERATE = 15  
SHAPE_DETECTION_THRESH = 0.7  
MIN_CONTOUR_AREA = 100 
DEBUG_MODE = True  

ser = None
try:
    ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=1)
    print(f"Serial port opened: {SERIAL_PORT}")
except Exception as e:
    print(f"Serial error: {e}")
    ser = None

dic_lables = {
    0: 'D',
    1: 'P',
    2: 'C',
    3: 'F',
    4: 'A',
    5: 'B'
}

def post_process_opencv(outputs, model_h, model_w, img_h, img_w, thred_nms, thred_cond):
    # outs: 模型输出，例如 (num_boxes, 6) -> [x1,y1,x2,y2,conf,cls]
    conf = outputs[:,4].tolist()
    c_x = outputs[:,0]/model_w * img_w
    c_y = outputs[:,1]/model_h * img_h
    w = outputs[:,2]/model_w * img_w
    h = outputs[:,2]/model_h * img_h
    p_cls = outputs[:,5:]
    if len(p_cls.shape)==1:
        p_cls = np.argmax(p_cls,axis = 1)
        
    p_x1 = np.expand_dims(c_x-w/2,-1)
    p_y1 = np.expand_dims(c_y-h/2,-1)
    p_x2 = np.expand_dims(c_x+w/2,-1)
    p_y2 = np.expand_dims(c_y+h/2,-1)
    areas = np.concatenate((p_x1,p_y1,p_x2,p_y2),axis = -1)
    
    areas = araes.tolist()
    ids = cv2.dnn.NMSBoxes(areas,conf,thred_cond,thred_nms)
    return np.array(areas)[ids],np.array(conf)[ids],cls_id[ids]

"""
def cal_outputs(outs,n1,na,model_w,model_h,anchor_grid,stride):
    
    row_ind = 0
    grid = [np.zeros(1)] * n1
    for i in range(n1):
        h,w = int(model_w/ stride[i]),int(model_h / stride[i])
        length = int(na * h * w)
        if grid[i].shape[2:4] != (h,w):
            grid[i] = _make_grid(w,h)
            
        outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
            grid[i],(na,1))) * int(stride[i])
        outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2 ) ** 2 * np.repeat(
            anchor_grid[i], h*w ,axis=0)
        row_ind += length  
    return outs
"""
def cal_outputs(outs, n1, na, model_w, model_h, anchor_grid, stride):
    print("outs.shape:", outs.shape)
    print("n1:", n1)
    print("na:", na)
    print("anchor_grid:", anchor_grid)
    print("stride:", stride)
    
    boxes = []
    scores = []
    ids = []

    batch_size, grid_h, grid_w, _ = outs.shape
    for i in range(batch_size):  # 遍历每个尺度
        for y in range(grid_h):
            for x in range(grid_w):
                for a in range(na):  # 遍历每个锚点
                    # 获取预测值
                    raw_pred = outs[i,y,x,:]
                    
                    # 提取预测结果
                    cx = (sigmoid(raw_pred[0]) * 2. - 0.5 + x) * stride[i]
                    cy = (sigmoid(raw_pred[1]) * 2. - 0.5 + y) * stride[i]
                    w = pow(sigmoid(raw_pred[2]) * 2, 2) * anchor_grid[i][a][0]
                    h = pow(sigmoid(raw_pred[3]) * 2, 2) * anchor_grid[i][a][1]
                    
                    # 计算边界框左上角坐标
                    x_min = cx - w / 2
                    y_min = cy - h / 2
                    x_max = cx + w / 2
                    y_max = cy + h / 2
                    
                    # 对象置信度
                    obj_conf = sigmoid(raw_pred[4])
                    
                    # 类别得分
                    cls_scores = softmax(raw_pred[5:])
                    class_id = np.argmax(cls_scores)
                    score = obj_conf * cls_scores[class_id]

                    if score > 0:  # 过滤掉低置信度的检测
                        boxes.append([x_min, y_min, x_max, y_max])
                        scores.append(score)
                        ids.append(class_id)

    return np.array(boxes), np.array(scores), np.array(ids)

# 辅助函数定义
def sigmoid(x):
    return 1. / (1. + np.exp(-x))

def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis = -1, keepdims = True)


def _make_grid(nx,ny):
    xv,yv = np.meshgrid(np.arange(ny),np.arange(nx))
    return np.stack((xv,yv),2).reshape((-1,2)).astype(np.float32)

def infer_image(net, img, model_h, model_w,n1,na,stride,anchor_grid, thred_nms=0.4, thred_cond=0.5):
    
    if img is None or img.size == 0:
        return np.array([]), np.array([]), np.array([])
    #yuchuli
    img0 = img.copy()
    img0 = cv2.resize(img0, [model_w, model_h],interpolation = cv2.INTER_AREA)
    img0 = cv2.cvtColor(img0,cv2.COLOR_BGR2RGB)
    img0 = img0.astype(np.float32) / 255.0
    blob = np.expand_dims(np.transpose(img0,(2,0,1)),axis = 0)
    
    #tuili
    outs = net.run(None,{net.get_inputs()[0].name:blob})[0].squeeze(axis=0)
    
    #zuobiaojiaozhe
    print("outs shape:", outs.shape)
    outs = cal_outputs(outs,n1,na,model_w,model_h,anchor_grid,stride)
    
    #jiancekuangjisuan
    img_h,img_w,_ = np.shape(img0)
    boxes,confs,ids = post_process_opencv(outs, model_h, model_w, img_h, img_w, thred_nms, thred_cond)
    
    return boxes,confs,ids
    
def detect_shape_simple(roi):
    
    try:
        
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
                
        _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)        
        
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return "unknown"
        
        cnt = max(contours, key=cv2.contourArea)
        
        if cv2.contourArea(cnt) < MIN_CONTOUR_AREA:
            return "unknown"
        
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.04 * peri, True)
        
        if len(approx) == 4:
            return "R"
        elif len(approx) > 8:  
            return "E"
        
    except:
        pass
    
    return "E"

def send_detections(detections):
    if ser is None or not ser.is_open or not detections:
        return

    try:
        for (label, confidence, box, shape) in detections:
            
            label_num = list(dic_lables.values()).index(label) if label in dic_lables.values() else 255
            shape_num = {"C": 1, "R": 2}.get(shape, 1)

            x = int(box[0])
            y = int(box[1])

            x = max(0, min(255, x))
            y = max(0, min(255, y))
            label_num = max(1, min(255, label_num))
            shape_num = max(1, min(255, shape_num))

            packet = bytes([
                0xFF,           
                x,             
                y,              
                label_num,      
                shape_num,      
                0xFE            
            ])

            ser.write(packet)

        if DEBUG_MODE:
            print(f"Sent {len(detections)} objects")

    except Exception as e:
        if DEBUG_MODE:
            print(f"Serial error: {e}")

def main():

    model_h = 320
    model_w = 320
    file_model = "/home/txw/picture_find/lite_trace.onnx"
    net = ort.InferenceSession(file_model)
    
    n1 = 3
    na = 3
    stride=[8.,16.,32.]
    anchors = [[10,13,16,30,33,23],[30,61,62,45,59,119],[116,90,156,198,373,326]]
    anchor_grid = np.asarray(anchors,dtype = np.float32).reshape(n1,-1,2)
    
    try:
        net = ort.InferenceSession(file_model)
        print(f"Model loaded")
    except Exception as e:
        print(f"Model error: {e}")
        return

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Failed to open camera")
        return
    
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, RESOLUTION[0])
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, RESOLUTION[1])
    cap.set(cv2.CAP_PROP_FPS, FRAMERATE)
    
    last_send_time = 0
    fps_counter = 0
    fps_timer = time.time()
    frame_counter = 0
    
    start_sending = False

    try:
        while True:
            ret, frame = cap.read()
            if not ret: 
                cap.release()
                time.sleep(0.5)
                cap = cv2.VideoCapture(0)
                if not cap.isOpened():
                    print("Camera reconnect failed")
                    break
                cap.set(cv2.CAP_PROP_FRAME_WIDTH, RESOLUTION[0])
                cap.set(cv2.CAP_PROP_FRAME_HEIGHT, RESOLUTION[1])
                cap.set(cv2.CAP_PROP_FPS, FRAMERATE)
                continue

            frame_counter += 1
            if frame_counter % 2 != 0:
                continue
                
            start_time = time.time()
            det_boxes, scores, ids = infer_image(net,frame, model_h, model_w,n1,na,stride,anchor_grid, thred_nms=0.4, thred_cond=0.5)
            inference_time = time.time() - start_time
            
            current_time = time.time()
            detections = []
            if len(ids) > 0:
                for box, score, id in zip(det_boxes, scores, ids):
                    
                    if score < SHAPE_DETECTION_THRESH:
                        continue
                        
                    label = dic_lables.get(int(id), "unknown")
                    
                    x1, y1, x2, y2 = [int(coord) for coord in box]
                  
                    x1, y1 = max(0, x1), max(0, y1)
                    x2, y2 = min(frame.shape[1]-1, x2), min(frame.shape[0]-1, y2)
                    
                    if x2 > x1 and y2 > y1 and (x2 - x1) > 10 and (y2 - y1) > 10:
                        roi = frame[y1:y2, x1:x2]
                        shape = detect_shape_simple(roi)
                    else:
                        shape = "E"
                    
                    detections.append((label, float(score), box, shape))
            
            if ser and ser.is_open:
                if ser.in_waiting > 0:
                    received_data = ser.read(ser.in_waiting)
                    if b'3' in received_data:
                        start_sending = True
                        print("Start sending triggered by host (received '3')")
                        
            if current_time - last_send_time >= SEND_INTERVAL and detections and start_sending:
                send_detections(detections)
                last_send_time = current_time
            
            fps_counter += 1
            if current_time - fps_timer >= 1.0:
                fps_value = fps_counter / (current_time - fps_timer)
                fps_counter = 0
                fps_timer = current_time
                print(f"FPS: {fps_value:.1f} | Inf: {inference_time*1000:.1f}ms | Obj: {len(detections)}")
            
            if DEBUG_MODE: 
                for (label, score, box, shape) in detections:
                    x1, y1, x2, y2 = [int(coord) for coord in box]
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f"{label}:{shape}:{score:.2f}", (x1, y1-10), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                cv2.imshow('Object Detection', frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break      
    finally:
        
        cap.release()
        if DEBUG_MODE:
            cv2.destroyAllWindows()
        if ser and ser.is_open:
            ser.close()
        print("Program exited.")

if __name__ == "__main__":
    main()